import platgo as pg
import numpy as np


class KP(pg.Problem):

    def __init__(self, D: int = 250, P: np.ndarray = None, W: np.ndarray = None):
        """
        :param D: dimension of decs
        :param P: Profit of each item 1-D array
        :param W: Weight of each item 1-D array
        """
        self.name = "KP"
        self.borders = []
        self.type['single'], self.type['binary'], self.type['large'], self.type['constrained'] = [True] * 4
        self.M = 1
        self.D = D
        # TODO add file operation
        if P is None or W is None:
            self.P = np.random.randint(low=10, high=100, size=(1, self.D))
            self.W = np.random.randint(low=10, high=100, size=(1, self.D))
        else:
            self.P = P
            self.W = W
            self.D = P.shape[0]
        super().__init__()

    def cal_obj(self, pop: pg.Population):
        pop.objv = sum(self.P) - np.dot(pop.decs, self.P.T)

    def cal_cv(self, pop: pg.Population):
        pop.cv = np.dot(pop.decs, self.W.T) - sum(self.W)/2

    def get_optimal(self) -> np.ndarray:
        pass


if __name__ == "__main__":
    p = np.array([1,2,3,4,9,8,7])
    w = np.array([1,2,3,4,5,6,7])
    pro = KP(P=p, W=w)
    print(pro.D)
    decs = np.array([[0,1,0,1,1,1,1],[1,0,0,0,1,1,1]])
    pop = pg.Population(decs=decs)
    pro.cal_obj(pop)
    print(pop.objv)
    pro.cal_cv(pop)
    print(pop.cv)


